from abc import ABC, abstractmethod
from contextlib import asynccontextmanager
from typing import List, Optional

from pydantic import BaseModel, Field, model_validator

from src.llm import LLM
from src.logger import logger
from src.schema import ROLE_TYPE, AgentState, Memory, Message


class BaseAgent(BaseModel, ABC):
    """智能代理基础抽象类，用于管理代理状态和执行流程。
    
    提供状态转换、内存管理和基于步骤的执行循环等基础功能。
    子类必须实现 `step` 方法来定义具体的执行逻辑。
    
    主要功能：
    - 状态管理：IDLE、RUNNING、FINISHED、ERROR 等状态转换
    - 内存管理：消息历史记录的存储和检索
    - 执行控制：最大步数限制、死循环检测
    - 安全执行：异常处理和状态恢复
    """

    # 核心属性
    name: str = Field(..., description="代理的唯一名称")
    description: Optional[str] = Field(None, description="代理的可选描述信息")

    # 提示词配置
    system_prompt: Optional[str] = Field(
        None, description="系统级指令提示词，定义代理的角色和行为"
    )
    next_step_prompt: Optional[str] = Field(
        None, description="下一步行动提示词，指导代理决策过程"
    )

    # 依赖组件
    llm: LLM = Field(default_factory=LLM, description="大语言模型实例，用于生成响应")
    memory: Memory = Field(default_factory=Memory, description="代理的记忆存储，保存对话历史")
    state: AgentState = Field(
        default=AgentState.IDLE, description="代理当前状态：空闲/运行/完成/错误"
    )

    # 执行控制参数
    max_steps: int = Field(default=10, description="最大执行步数，防止无限循环")
    current_step: int = Field(default=0, description="当前执行步数计数器")

    # 死循环检测阈值：连续重复响应的次数上限
    duplicate_threshold: int = 2

    class Config:
        """Pydantic 模型配置类。"""
        # 允许任意类型，支持复杂对象如 LLM 实例
        arbitrary_types_allowed = True
        # 允许额外字段，为子类提供灵活性
        extra = "allow"

    @model_validator(mode="after")
    def initialize_agent(self) -> "BaseAgent":
        """初始化代理，如果未提供则使用默认设置。
        
        确保 LLM 和 Memory 组件正确初始化，如果没有提供有效实例，
        则创建默认配置的组件。
        """
        # 初始化 LLM 实例，使用代理名称作为配置名
        if self.llm is None or not isinstance(self.llm, LLM):
            self.llm = LLM(config_name=self.name.lower())
        
        # 初始化内存组件
        if not isinstance(self.memory, Memory):
            self.memory = Memory()
        
        return self

    @asynccontextmanager
    async def state_context(self, new_state: AgentState):
        """代理状态安全转换的上下文管理器。
        
        提供代理状态的安全切换机制，确保在执行过程中状态的正确管理，
        即使发生异常也能正确恢复到之前的状态。

        Args:
            new_state: 要转换到的新状态

        Yields:
            None: 允许在新状态下执行代码

        Raises:
            ValueError: 如果提供的新状态无效
        """
        # 验证状态的有效性
        if not isinstance(new_state, AgentState):
            raise ValueError(f"无效的状态类型: {new_state}")

        # 保存当前状态以便恢复
        previous_state = self.state
        # 切换到新状态
        self.state = new_state
        
        try:
            yield
        except Exception as e:
            # 发生异常时，将状态设置为错误状态
            self.state = AgentState.ERROR
            logger.error(f"代理 {self.name} 执行异常，状态已设置为ERROR: {str(e)}")
            raise e
        finally:
            # 无论是否发生异常，都恢复到之前的状态
            # 除非当前状态是 ERROR，则保持 ERROR 状态
            if self.state != AgentState.ERROR:
                self.state = previous_state

    def update_memory(
        self,
        role: ROLE_TYPE,  # type: ignore
        content: str,
        base64_image: Optional[str] = None,
        **kwargs,
    ) -> None:
        """向代理的内存中添加消息。
        
        支持多种角色的消息类型，包括用户、系统、助手和工具消息。
        自动根据角色类型创建相应的消息对象并存储到内存中。

        Args:
            role: 消息发送者的角色类型 (user/system/assistant/tool)
            content: 消息内容文本
            base64_image: 可选的 base64 编码图像数据
            **kwargs: 额外参数 (如工具消息的 tool_call_id)

        Raises:
            ValueError: 如果提供的角色类型不被支持
        """
        # 消息类型映射表，根据角色创建对应的消息对象
        message_map = {
            "user": Message.user_message,
            "system": Message.system_message, 
            "assistant": Message.assistant_message,
            "tool": lambda content, **kw: Message.tool_message(content, **kw),
        }

        # 验证角色类型是否支持
        if role not in message_map:
            raise ValueError(f"不支持的消息角色类型: {role}")

        # 根据角色类型构建参数，工具消息需要额外的参数
        message_kwargs = {
            "base64_image": base64_image, 
            **(kwargs if role == "tool" else {})
        }
        
        # 创建并添加消息到内存
        message = message_map[role](content, **message_kwargs)
        self.memory.add_message(message)
        
        logger.debug(f"代理 {self.name} 添加 {role} 消息到内存: {content[:50]}...")

    async def run(self, request: Optional[str] = None) -> str:
        """异步执行代理的主循环。
        
        代理的核心执行方法，负责管理整个任务执行流程，包括：
        1. 状态检查和切换
        2. 请求处理和内存管理
        3. 循环执行步骤和结果记录
        4. 死循环检测和处理

        Args:
            request: 可选的初始用户请求，将被添加到对话历史中

        Returns:
            执行结果的文本总结，包含每个步骤的详细信息

        Raises:
            RuntimeError: 如果代理不在 IDLE 状态时尝试执行
        """
        # 检查代理状态，只有在 IDLE 状态才能开始执行
        if self.state != AgentState.IDLE:
            raise RuntimeError(f"无法从状态 {self.state} 启动代理，必须在 IDLE 状态")

        # 如果提供了请求，将其添加到对话历史
        if request:
            self.update_memory("user", request)
            logger.info(f"代理 {self.name} 开始处理用户请求: {request[:100]}...")

        execution_results: List[str] = []
        
        # 使用上下文管理器安全切换到 RUNNING 状态
        async with self.state_context(AgentState.RUNNING):
            # 主执行循环：在最大步数限制内且未完成时继续执行
            while (
                self.current_step < self.max_steps and 
                self.state != AgentState.FINISHED
            ):
                self.current_step += 1
                logger.info(f"代理 {self.name} 正在执行第 {self.current_step}/{self.max_steps} 步")
                
                # 执行单个步骤（由子类实现）
                step_result = await self.step()

                # 检测是否陷入死循环并处理
                if self.is_stuck():
                    logger.warning(f"代理 {self.name} 检测到死循环，正在处理...")
                    self.handle_stuck_state()

                # 记录步骤执行结果
                execution_results.append(f"步骤 {self.current_step}: {step_result}")

            # 检查是否达到最大步数限制
            if self.current_step >= self.max_steps:
                self.current_step = 0
                self.state = AgentState.IDLE
                timeout_msg = f"执行终止：达到最大步数限制 ({self.max_steps})"
                execution_results.append(timeout_msg)
                logger.warning(f"代理 {self.name} {timeout_msg}")
        
        # 返回执行结果汇总
        result_summary = "\n".join(execution_results) if execution_results else "未执行任何步骤"
        logger.info(f"代理 {self.name} 执行完成，共 {len(execution_results)} 个步骤")
        return result_summary

    @abstractmethod
    async def step(self) -> str:
        """执行代理工作流中的单个步骤。
        
        这是一个抽象方法，必须由子类实现来定义具体的执行逻辑。
        每个步骤可能包括：
        - 分析当前情况和上下文
        - 决定下一步行动
        - 执行工具调用或生成响应
        - 更新内部状态
        
        Returns:
            步骤执行结果的描述文本
        """

    def handle_stuck_state(self) -> None:
        """处理代理陷入死循环的情况。
        
        当检测到代理重复生成相同的响应时，通过添加提示词
        来引导代理尝试不同的策略和方法，避免继续陷入死循环。
        """
        # 死循环处理提示词
        stuck_prompt = (
            "检测到重复响应。请考虑采用新的策略，避免重复已经尝试过的无效路径。"
            "可以尝试：1) 改变分析角度 2) 使用不同的工具 3) 给出结论性回答"
        )
        
        # 将提示词添加到下一步提示词的开头
        if self.next_step_prompt:
            self.next_step_prompt = f"{stuck_prompt}\n\n{self.next_step_prompt}"
        else:
            self.next_step_prompt = stuck_prompt
            
        logger.warning(f"代理 {self.name} 检测到死循环，已添加引导提示")

    def is_stuck(self) -> bool:
        """检测代理是否陷入死循环，通过识别重复的响应内容。
        
        通过对比最近的助手消息与历史消息，统计相同内容出现的次数。
        如果连续重复次数超过阈值，则认为代理陷入了死循环。
        
        Returns:
            True 如果检测到死循环，False 否则
        """
        # 检查消息数量是否足够进行对比
        if len(self.memory.messages) < 2:
            return False

        # 获取最后一条消息
        last_message = self.memory.messages[-1]
        
        # 检查消息内容是否为空或无效
        if not last_message.content or not last_message.content.strip():
            return False

        # 统计相同内容在历史中的出现次数（只考虑助手消息）
        duplicate_count = 0
        target_content = last_message.content.strip()
        
        # 倒序遍历历史消息（排除最后一条）
        for msg in reversed(self.memory.messages[:-1]):
            if (
                msg.role == "assistant" and 
                msg.content and 
                msg.content.strip() == target_content
            ):
                duplicate_count += 1

        # 判断是否超过重复阈值
        is_stuck = duplicate_count >= self.duplicate_threshold
        
        if is_stuck:
            logger.debug(f"代理 {self.name} 检测到死循环：相同内容重复 {duplicate_count} 次")
            
        return is_stuck

    @property
    def messages(self) -> List[Message]:
        """获取代理内存中的消息列表。
        
        提供对代理对话历史的只读访问，方便外部代码查看对话内容。
        
        Returns:
            包含所有对话消息的列表
        """
        return self.memory.messages

    @messages.setter
    def messages(self, value: List[Message]) -> None:
        """设置代理内存中的消息列表。
        
        允许外部代码更新代理的对话历史，通常用于：
        - 恢复之前的对话状态
        - 清空对话历史
        - 导入外部对话数据
        
        Args:
            value: 新的消息列表
        """
        self.memory.messages = value
        logger.debug(f"代理 {self.name} 对话历史已更新，包含 {len(value)} 条消息")

    def reset_execution_state(self) -> None:
        """重置代理的执行状态，为新一轮执行做准备。
        
        该方法会：
        1. 将步数计数器重置为 0
        2. 将代理状态返回到 IDLE
        3. 保留对话历史和其他配置
        
        通常在以下情况下使用：
        - 上一次执行出现错误后重新开始
        - 在多轮对话中重置执行计数
        - 清理执行状态但保留对话历史
        """
        previous_step = self.current_step
        previous_state = self.state
        
        # 重置执行参数
        self.current_step = 0
        self.state = AgentState.IDLE
        
        logger.info(
            f"代理 '{self.name}' 执行状态已重置: "
            f"步数 {previous_step} -> 0, 状态 {previous_state} -> {AgentState.IDLE}"
        )
